import abc

import numpy as np
import modern_robotics as mr


class Robot(abc.ABC):

    def __init__(self):
        super().__init__()
        self._dof = 0
        self._q = np.zeros(self._dof)

        self._M = np.eye(4)
        self._Tgb = np.eye(4)
        self._S_array = np.zeros((6, self._dof))
        self._B_array = np.zeros((6, self._dof))
        self._Ms = []
        self._q_lim = np.zeros((2, self._dof))
        self._qd_lim = np.zeros(self._dof)

    @property
    def dof(self):
        return self._dof

    @property
    def q(self):
        return self._q.copy()

    @q.setter
    def q(self, q: np.ndarray):
        self._q[:] = q

    @property
    def qd_lim(self):
        return self._qd_lim.copy()

    def fkine(self, q: np.ndarray, include_base=True):
        T = mr.FKinSpace(self._M, self._S_array, q)
        if not include_base:
            T = mr.TransInv(self._Tgb) @ T
        return T

    def jacob0(self, q: np.ndarray, start=0, end=-1):
        q_end = q.copy()
        if end < 0:
            end = len(self._Ms) - 1
        q_end[end:] = 0
        T_end = mr.FKinSpace(self._Tgb @ self._Ms[end], self._S_array, q_end)
        JacobianSpace = mr.JacobianSpace(self._S_array, q)
        jacob0 = np.zeros_like(JacobianSpace)
        jacob0[3:, :] = JacobianSpace[:3, :]
        for i in range(self._dof):
            jacob0[:3, i] = JacobianSpace[3:, i] + np.cross(JacobianSpace[:3, i], T_end[:3, 3])

        q_start = q.copy()
        q_start[start:] = 0
        T_start = mr.FKinSpace(self._Tgb @ self._Ms[start], self._S_array, q_start)
        RR = np.eye(6)
        RR[:3, :3] = T_start[:3, :3]
        RR[3:6, 3:6] = RR[:3, :3]
        return RR.T @ jacob0[:, start:]

    def jacobe(self, q: np.ndarray):
        JacobianBody = mr.JacobianBody(self._B_array, q)
        jacobe = np.zeros_like(JacobianBody)
        jacobe[:3, :] = JacobianBody[3:, :]
        jacobe[3:, :] = JacobianBody[:3, :]
        return jacobe

    def hessian0(self, q: np.ndarray, start=0, end=-1, J=None):
        H = np.zeros((self._dof - start, 6, self._dof - start))
        if J is None:
            J = self.jacob0(q, start=start, end=end)
        for i in range(self._dof - start):
            for j in range(self._dof - start):
                a = np.min([i, j])
                b = np.max([i, j])
                H[i, :3, j] = np.cross(J[3:, a], J[:3, b])
                if i < j:
                    H[i, 3:, j] = np.cross(J[3:, i], J[3:, j])
        return H

    def manipulability(self, q, start=0, end=-1, J=None):
        if J is None:
            J = self.jacob0(q, start=start, end=end)
        return np.sqrt(np.abs(np.linalg.det(J @ J.T)))

    def jacobm(self, q, start=0, end=-1, J=None, H=None):
        if J is None:
            J = self.jacob0(q, start=start, end=end)
        n = J.shape[1]
        if H is None:
            H = self.hessian0(q, start=start, end=end, J=J)

        manipulability = self.manipulability(q, start=start, end=end, J=J)
        b = np.linalg.pinv(J @ J.T)
        Jm = np.zeros((n, 1))

        for i in range(n):
            c = J @ np.transpose(H[i, :, :])
            Jm[i, 0] = manipulability * np.transpose(c.flatten("F")) @ b.flatten("F")

        return Jm

    def joint_velocity_damper(self, q, ps, pi, gain=1.0):
        n = q.size
        Ain = np.zeros((n, n))
        Bin = np.zeros(n)

        for i in range(n):
            if q[i] - self._q_lim[0, i] <= pi:
                Bin[i] = -gain * ((self._q_lim[0, i] - q[i]) + ps) / (pi - ps)
                Ain[i, i] = -1
            if self._q_lim[1, i] - q[i] <= pi:
                Bin[i] = gain * ((self._q_lim[1, i] - q[i]) - ps) / (pi - ps)
                Ain[i, i] = 1

        return Ain, Bin

    def p_servo(self, T_actual, T_desired, gain=1.0, threshold=0.1):
        Vb = np.zeros(6)
        T_err = mr.TransInv(T_actual) @ T_desired
        Vb[:3] = T_err[:3, -1]
        Vb[3:] = mr.so3ToVec(mr.MatrixLog3(T_err[:3, :3]))
        v = gain * Vb
        arrived = True if np.sum(np.abs(Vb)) < threshold else False
        return v, arrived

    def set_base(self, T_base: np.ndarray):
        self._Tgb[:, :] = T_base
        self._M[:, :] = T_base @ self._Tbe @ self._Tet
        self._S_array[:, :] = mr.Adjoint(self._Tgb) @ np.vstack((self._w_array, self._v_array))
        self._B_array[:, :] = mr.Adjoint(mr.TransInv(self._M)) @ self._S_array
